import struct
import socket

class Socks5Proxy(socket.socket):
	''' Mimic a normal socket '''
	def __init__(self, host, port, timeout=10.0):
		self.host = host
		self.port = port
		socket.socket.__init__( self, socket.AF_INET, socket.SOCK_STREAM )
		self.settimeout(timeout)
		
	def proxy_connect(self, tserver, userName = None, password = None):
		thost, tport = tserver
		print "[PROXY] Connecting to proxy %s:%s" % (self.host, self.port)
		self.connect( (self.host, self.port) )
		print "[PROXY] Connected!"
		self.settimeout(0)
		buf = chr(0x05)
		if userName is not None:
			buf += chr(0x02) + chr(0x00) + chr(0x02)
		else:
			buf += chr(0x01) + chr(0x00)

		self.sendall(buf)

		buf = self.recv(2)
		if ord(buf[0]) != 0x05 or ord(buf[1]) == 0xff:
			print "[PROXY] SOCKS5 Error: Proxy gave a bad response to connection information"
			return False

		if userName is not None:
			buf = '\x01' + struct.pack('!B', len(userName)) + userName
			buf += struct.pack('!B', len(password)) + password

			self.sendall(buf)

			buf = self.recv(2)
			if len(buf) != 2 or ord(buf[0]) != 0x01 or ord(buf[1]) != 0x00:
				print "[PROXY] SOCKS5 Error: Proxy gave a bad answer to username"
				return False

		buf = '\x05\x01\x00\x03'
		buf += struct.pack('!B', len(thost)) + thost
		buf += struct.pack('!H', tport)
		self.send(buf)

		buf = self.recv(100)
		if ord(buf[0]) != 0x05 or ord(buf[1]) != 0x00: return "SOCKS5 Error: Proxy gave a bad response to host information"
		print "[PROXY] Authorized!"
		return True
